# 导包
import os
import time
import tensorflow as tf
import pandas as pd
import matplotlib.pyplot as plt
import sys
sys.path[0] = '../../imagenette2-160/'

from tensorflow.keras.models import save_model
from nets.conv_net import ConvModel
from utils.data_generator import train_val_generator
from utils.image_plot import plot_image



# 读取训练集
train_gen = train_val_generator(

    data_dir='../dataset/imagenette2-160/train',
    target_size=(160, 160),
    batch_size=32,
    class_mode='categorical',
    subset='training'
)
# 读取验证集
val_gen = train_val_generator(
    data_dir='../dataset/imagenette2-160/train',
    target_size=(160, 160),
    batch_size=32,
    class_mode='categorical',
    subset='validation'
)

# 获取class_names的值
class_names = list(train_gen.class_indices.keys())
print(class_names)

# ImageDataGenerator的返回结果是个迭代器，调用一次才会吐一次结果
# 可以使用.next()函数分读取图片
train_batch, train_label_batch = train_gen.next()

"""
plot_image(train_batch, train_label_batch)
"""
plot_image(train_batch, train_label_batch, class_names)

# 取15
val_batch, val_label_batch = val_gen.next()

"""
plot_image(val_batch, val_label_batch)
"""
plot_image(val_batch, val_label_batch, class_names)

# 类实例化
model = ConvModel()

# 用到的参数
# 设置损失函数loss,优化器optimizer,评价标准metrics
model.compile(
    loss='categorical_crossentropy',
    optimizer=tf.keras.optimizers.SGD(
        learning_rate=0.001),
    metrics=['accuracy'])

"""
模型训练tf.keras.Sequential.fit
- x: 输入的训练集，可以用ImageDataGenerator读取的数据
- steps_per_epoch: 输入整数，数据集跑多少轮
- validation_data
_ validation_steps: 验证集跑多少步来计算模型的评价指标 一步读取batch_size张图片，
- 所以一共验证validation_steps*batch_size
- shuffle: 每轮训练是否会打乱孙旭，默认ture
返回：
- History 对象：记录每一轮训练集和验证集的损失函数和评价指标

"""

# 模型训练
history = model.fit(x=train_gen, steps_per_epoch=235,
                    validation_data=val_gen, epochs=100,
                    shuffle=True)
# 画图查看history数据的变化趋势
pd.DataFrame(history.history).plot(figsize=(8, 5))
plt.grid(True)
plt.xlabel('epoch')
plt.show()

# 创建保持路径
# model_name = "model-" + time.strfitime('%Y-%m-%d-%H-%M-%S')
model_name = "model-AlX"
model_path = os.path.join('..', 'model', model_name)
if not os.path.exists(model_path):
    os.makedirs(model_path)

# 模型保存
save_model(model=model, filepath=model_path)
